package org.dreamwork.sso.saml2.filter;

import org.dreamwork.sso.saml2.SamlKeys;
import org.dreamwork.sso.saml2.protocol.SamlLoginRequest;
import org.dreamwork.sso.saml2.protocol.SamlStatus;
import org.dreamwork.sso.saml2.spec.IKeyFetcher;
import org.dreamwork.sso.saml2.spec.KeyFetcherFactory;
import org.dreamwork.secure.AlgorithmMapping;
import org.dreamwork.sso.saml2.util.ConfigKey;
import org.dreamwork.misc.PathFilter;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.dreamwork.util.StringUtil;
import org.joda.time.DateTime;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.IOException;
import java.security.PrivateKey;
import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * User: seth.yang
 * Date: 13-8-22
 * Time: 下午5:37
 */
public class SsoLoginFilter implements Filter {
    private Log logger = LogFactory.getLog (SsoLoginFilter.class);

    private Set<PathFilter> excludes = new HashSet<PathFilter> ();

    @Override
    public void init (FilterConfig config) throws ServletException {
        String expression = config.getInitParameter ("exclude");
        if (!StringUtil.isEmpty (expression)) {
            String[] parts = expression.split ("[,\\s\\|]");
            for (String part : parts)
                excludes.add (new PathFilter (part.trim ()));
        }
    }

    @Override
    public void doFilter (ServletRequest req, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        HttpServletRequest request = (HttpServletRequest) req;
        if (!isExclude (request)) {
            HttpSession session = request.getSession ();
            SamlStatus status = (SamlStatus) session.getAttribute (SamlKeys.SAML_LOGIN_STATUS);
            if (status == null) {
                backupCurrentRequest (request);
                redirectToIdp (request, (HttpServletResponse) response);
                return;
            }
        }

        chain.doFilter (request, response);
    }

    @Override
    public void destroy () {
        excludes.clear ();
    }

    private boolean isExclude (HttpServletRequest request) {
        String contextPath = request.getContextPath ();
        String uri = request.getRequestURI ();
        uri = uri.substring (contextPath.length ());
        for (PathFilter filter : excludes) {
            if (filter.hit (uri)) return true;
        }
        return false;
    }

    private void redirectToIdp (HttpServletRequest request, HttpServletResponse response) throws IOException {
        HttpSession session = request.getSession ();
        ServletContext context = session.getServletContext ();

        Properties config = (Properties) context.getAttribute (SamlKeys.CONFIG_CONTEXT_KEY);

        SamlLoginRequest samlLoginRequest = new SamlLoginRequest ();
        samlLoginRequest.setAssertionConsumerServiceURL (config.getProperty (ConfigKey.SAML2_ASSERTION_CONSUMER_URL));
        samlLoginRequest.setDestination (config.getProperty (ConfigKey.SAML2_IDP_URL));
        samlLoginRequest.setIssuer (config.getProperty (ConfigKey.SAML2_ISSUER));
        samlLoginRequest.setId ("_" + UUID.randomUUID ().toString ());
        samlLoginRequest.setIssueInstant (new DateTime ());

        KeyFetcherFactory factory = KeyFetcherFactory.newInstance ();
        IKeyFetcher fetcher = factory.getFetcher ();
        PrivateKey key = fetcher.getPrivateKey (context, null);
        if (logger.isDebugEnabled ())
            logger.debug ("trying to sign the SAML2.0 Login Request ... ");
        String signature = samlLoginRequest.sign (key, AlgorithmMapping.Signature.RSA_SHA1.uri);
        String rawText = samlLoginRequest.createRawText (AlgorithmMapping.Signature.RSA_SHA1.uri);

        String url = String.format ("%s?%s&Signature=%s",
                config.getProperty (ConfigKey.SAML2_IDP_URL),
                rawText, signature
        );
        if (logger.isDebugEnabled ())
            logger.debug ("redirect to idp: " + url);
        response.sendRedirect (url);
    }

    @SuppressWarnings ("unchecked")
    private void backupCurrentRequest (HttpServletRequest request) {
        HttpSession session = request.getSession ();
        Map<String, Object> backup = new WeakHashMap<String, Object> ();
        String method = request.getMethod ();
        backup.put ("method", method);
        backup.put ("url", request.getRequestURL ());
        if ("GET".equalsIgnoreCase (method)) {
            String queryString = request.getQueryString ();
            if (!StringUtil.isEmpty (queryString))
                backup.put ("queryString", request.getQueryString ());
        } else if ("POST".equalsIgnoreCase (method)) {
            Enumeration<String> e = request.getParameterNames ();
            Map<String, Object> parameters = new HashMap<String, Object> ();
            while (e.hasMoreElements ()) {
                String name = e.nextElement ();
                String[] values = request.getParameterValues (name);
                if (values.length == 0)
                    parameters.put (name, "");
                else if (values.length == 1)
                    parameters.put (name, values [0]);
                else
                    parameters.put (name, values);
            }
            if (parameters.size () > 0)
                backup.put ("parameters", parameters);
        }
        session.setAttribute (SamlKeys.SAML_REQUEST_BACKUP, backup);
    }
}